#!/usr/bin/env python

import optparse
import sys

import minuit
import numpy
import tables

from charmicat import BadCommandLineError


def bg_log_llh(p0, p1, p2, p3, p4, p5, p6, p7):
    global _g_data, _g_mc, _g_ncalls

    llh = -numpy.sum(numpy.log(bg_intensity(_g_data, p0, p1, p2, p3, p4, p5, p6, p7)))

    norm_int = numpy.sum(bg_intensity(_g_mc, p0, p1 ,p2, p3, p4, p5, p6, p7))  

    norm_int *= (len(_g_data) / len(_g_mc))

    if _g_ncalls % 100 == 0:
        print "bg_log_llh", _g_ncalls, llh, norm_int, llh + norm_int

    _g_ncalls += 1

    return llh + norm_int

def bg_intensity(x, p0, p1, p2, p3, p4, p5, p6, p7):
    return numpy.polynomial.chebyshev.chebval(x, [p0, p1, p2, p3, p4, p5, p6, p7])



def main():
    '''
    ypipi_pi0pi0_llhfit.py [opts] file.h5 [file.h5 [...]]

    Fit histogram `mm_off_pi0pi0' from 9.35 to 9.7 GeV with Gauss + PolN
    using log-likelihood method.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    opts, args = op.parse_args()

    if len(args) < 1:
        msg = 'Please supply an input file.'
        raise BadCommandLineError(msg, -1)

    lower_limit        = 9.3
    signal_lower_limit = 9.4
    signal_upper_limit = 9.6
    upper_limit        = 9.7

    cut_str = '((mm_off_pi0pi0 > 9.35) & (mm_off_pi0pi0 < 9.4)) | ((mm_off_pi0pi0 > 9.6) & (mm_off_pi0pi0 < 9.7))'

    data_by_file = []
    for path in args:
        with tables.openFile(path) as fp:
            data_by_file.append(fp.root.nt.readWhere(cut_str, field='mm_off_pi0pi0'))

    global _g_data, _g_mc, _g_ncalls

    _g_data = numpy.concatenate(data_by_file)
    _g_mc   = numpy.concatenate((numpy.random.uniform(lower_limit, signal_lower_limit, len(_g_data)),
                                numpy.random.uniform(signal_upper_limit, upper_limit, len(_g_data))))

    m_bg = minuit.Minuit(bg_log_llh)
    
    m_bg.values['p0'] = 1
    m_bg.values['p1'] = 1
    m_bg.values['p2'] = 1
    m_bg.values['p3'] = 0
    m_bg.values['p4'] = 0
    m_bg.values['p5'] = 0
    m_bg.values['p6'] = 0

    m_bg.fixed['p4'] = True
    m_bg.fixed['p5'] = True
    m_bg.fixed['p6'] = True

    initial_values = {}
    initial_values.update(m_bg.values)
    
    _g_ncalls = 0

    try:
        m_bg.migrad()
    except minuit.MinuitError:
        raise

    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.6, 0.8))

    info = [m_bg.values['p0'], m_bg.errors['p0'],
            m_bg.values['p1'], m_bg.errors['p1'],
            m_bg.values['p2'], m_bg.errors['p2'],
            m_bg.values['p3'], m_bg.errors['p3'],
            m_bg.values['p4'], m_bg.errors['p4'],
            m_bg.values['p5'], m_bg.errors['p5'],
            m_bg.values['p6'], m_bg.errors['p6'],
            bg_chisq(**initial_values) / (len(_g_x) - len(m_bg.values.keys())), 
            bg_chisq(**m_bg.values) / (len(_g_x) - len(m_bg.values.keys())), 
            ]

    fptex = ypipi.textable.write_minuit_table(m_bg)

    

    pyplot.text(1.35, 0.75,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = bg(xfit, **m_bg.values)

    pyplot.plot(xfit, yfit, 'b-')
#    pyplot.plot(xfit, y_calc(xfit, **initial_values), 'k--')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title + r' [all4g $>$ 50 MeV]')

    path_out = h.name + '_bgfit.png'
    pyplot.savefig(path_out)
    



    
if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
